import pandas as pd
import numpy as np
import os
import info
import StrategyModule
import matplotlib.pyplot as plt
import datetime
from EnvironmentModule import Environment
from TradeRecordModule import TradeRecord
from EvaluationModule import Evaluation
import time

class mystrategy(StrategyModule.Strategy):

    def strategy_calculate(
            self,
                           trade_data):
        ###### 需要修改部分 ######
        close = trade_data.close
        ma_5 = close.rolling(5).mean()[-1]  # 最新均线值, T-1日收盘后数据
        ma_10 = close.rolling(10).mean()[-1]  # 最新均线值, T-1日收盘后数据
        ma_20 = close.rolling(20).mean()[-1]  # 最新均线值, T-1日收盘后数据

        self.enter_long_postion_signal = ma_20 < ma_10  # 开多信号
        self.enter_short_postion_signal = ma_20 > ma_10  # 开空信号
        self.exit_long_postion_signal = ma_10 < ma_5  # 平多信号
        self.exit_short_postion_signal = ma_10 > ma_5  # 平空信号
        ###### 需要修改部分 ######

    def trade_record(
                    self,
                    tradeRecord,
                    current_bar_contract_name,
                    current_bar_datetime,
                    close_position_price,
                    open_position_price,
                    last_bar_price,
                    daily_clearing_position_price,
                    clearing_daily,
                    freq,
                    current_bar_trade_date
                     ):
        # checkPosition
        tradeRecord.set_current_bar_trade_info(
                                                contract_name=current_bar_contract_name,
                                                close_position_price=close_position_price,
                                                open_position_price=open_position_price,
                                                last_bar_price=last_bar_price,
                                                current_bar_datetime=current_bar_datetime,
                                                daily_clearing_position_price=daily_clearing_position_price,
                                                current_bar_trade_date=current_bar_trade_date
        )
        tradeRecord.update_current_position()

        ###### 需要修改部分 ######
        # 开多
        if self.enter_long_postion_signal:
            tradeRecord.enter_long_position(
                # trade_unit=1,
                # position='long',
            )
        # 开空
        elif self.enter_short_postion_signal:
            # 平空
            tradeRecord.enter_short_position(
                # trade_unit=1,
                # position='short',
            )
        # 平多
        if self.exit_long_postion_signal:
            tradeRecord.exit_long_position()
        # 平空
        if self.exit_short_postion_signal:
            tradeRecord.exit_short_position()
        ###### 需要修改部分 ######

        if clearing_daily and ('D' in freq):
            tradeRecord.exit_long_position_clearing_daily(daily_clearing_position_price)
            tradeRecord.exit_short_position_clearing_daily(daily_clearing_position_price)

if __name__ == '__main__':
    end_date = '2022-06-24'
    env = Environment(
        data_name='中证500指数_5min',  # 回测数据名称 日频D, 秒S, 分钟min
        data_address='hfBar_ZL',  # 回测数据存放位置 高频为 hfBar_ZL, 日频为 dailyBar_ZL
        data_location='remote',  # local: 本地数据, remote: 服务器数据
        trade_price_name='yclose',  # 当前bar交易价格基准, 可改成其他字段
        daily_clearing_price_name='close',  # 当前bar每日强制平仓交易价格基准, 可改成其他字段, 针对日内有效
        max_lookback_bars=20,  # 最大回看bar数量, 根据策略不同需要改变
        max_trade_times=20,  # 每日最大交易次数
        trade_unit=1,  # 每次交易手数
        stop_win=0.03,  # 止盈 费前收益超过3%
        stop_loss=-0.03,  # 止损 费前收益低于-3%
        x_min_price=3, # 变动几个tick, 实际买入价格 = 交易价格基准 + x_min_price * tick_size, 实际卖出价格 = 交易价格基准 - x_min_price * tick_size
        long_freeze_period=900,  # 买入冷却时间, 平多后多长时间不再开多, 单位:秒
        short_freeze_period=900,  # 卖出冷却时间, 平空后多长时间不再开空, 单位:秒

    )
    env.backtest_setting(
        start_date='2022-05-01',  # 回测开始日期
        end_date='2022-06-24',   # 回测结束日期
        clearing_time=[],  # 强制平仓时间段, 在该时间段内头寸会被强行平掉, 且不会开新仓

        )
    s_name = 'MA_5_10_20'  # 策略名称
    s = mystrategy(strategy_name=s_name)  # 实例化策略
    env.add_strategy(s)  # 向回测环境加入策略
    start = time.time()
    env.backtest_running()  # 开始回测
    end = time.time()
    print(end-start)
    trade_record = env.trade_recorder.get_final_position(end_date)  # 获得全部交易记录
    e = Evaluation(product_name='中证500指数')  # 初始化绩效统计模块
    e.plot_record(trade_record)  # 生成回测结果
